function [prob_ini,theta_fs] = FirstStage(Action,AgeState,States,pproc)

% Two separate estimations: in sugar and not in sugar states

AuxReg = States.RegionsForExoStates*ones(1,size(AgeState,2));
for r=unique(States.RegionsForExoStates)';
    AverageOut(r) = (pproc.outside{r}.Cte/(1 - pproc.outside{r}.AR));
end


%% In sugar:
%SelInSugar = (AgeState>1 & ~isnan(States.SingleStateIdx) );
SelInSugar = (AgeState>1 & States.SingleStateIdx>0 );
A  = Action(SelInSugar);
S  = States.SingleStateIdx(SelInSugar);
OS = AuxReg(SelInSugar);

% Set covariates for logit:

OutReturn = AverageOut(OS)'; % Other crops return
Fixed = States.AllStatesMap(S,3);
Distance = States.MidPointsFixedStates(States.FixedStatesMap(Fixed,1),1);
n = length(S);
J = size(States.MidPointsFixedStates,2)-1;
FixedCov = zeros(n,J);

for k=1:J
    FixedCov(:,k) = States.MidPointsFixedStates(States.FixedStatesMap(Fixed,1+k),1+k);
end
Time  = States.AllStatesMap(S,2);
Age   = double(States.AllStatesMap(S,1));

% Covariates are action specific: abandon, do nothing and replant: 
X1 = [zeros(n,1) ones(n,1) zeros(n,1) Distance OutReturn zeros(n,J+1)];
X2 = zeros(size(X1));
X3 = [ones(n,1) zeros(n,1) Distance zeros(n,2) FixedCov Age];
X = [X1 X2 X3];

[theta_fs.b_suc] = clogit(A,X);

%% Not in sugar:
%SelNotInSugar = (AgeState==1 & ~isnan(States.SingleStateIdx) );
SelNotInSugar = (AgeState==1 & States.SingleStateIdx>0 );
A = Action(SelNotInSugar);
S = States.SingleStateIdx(SelNotInSugar);
OS = AuxReg(SelNotInSugar);

% Set covariates for logit:
OutReturn = AverageOut(OS)'; % Other crops return
Fixed = States.AllStatesMap(S,3);
Distance = States.MidPointsFixedStates(States.FixedStatesMap(Fixed,1),1);
n = length(S);
J = size(States.MidPointsFixedStates,2)-1;
FixedCov = zeros(n,J);

for k=1:J
    FixedCov(:,k) = States.MidPointsFixedStates(States.FixedStatesMap(Fixed,1+k),1+k);
end

% LU class:
LUclass = double(States.FixedStatesMap(Fixed, end)); % Get Other Land Use class state
LUclass = dummyvar(LUclass); % create dummy variables for clogit
LUclass = LUclass(:,2:end); % clogit will be run with constant, so eliminate first column (not pasture nor cropland)
% Sugarcane Road Distance (SRD)
SRD     = States.SRDMidPoints(States.AllStatesMap(S,4),1);

% Covariates are action specific: plant sugar and do nothing: 
X2 = [ones(n,1) zeros(n,1) Distance log(SRD) FixedCov LUclass];
X1 = [zeros(n,1) OutReturn zeros(n,2+J+2)];
X = [X1 X2];

[theta_fs.b_nsuc] = clogit(A,X);


%% Generate choice probabilities:
%prob_ini = zeros(size(States.AllStatesMap,1),3);


for r = unique(States.RegionsForExoStates)';
    prob_ini{r} = zeros(size(States.AllStatesMap,1),3); % save space for choice prob
    
    % First create choice probability for in sugar states:
    
    SugarIdxR   = (States.AllStatesMap(:,1) >= 2);
    SugarAllStates  = States.AllStatesMap(SugarIdxR,:);
    Distance = States.MidPointsFixedStates(States.FixedStatesMap(SugarAllStates(:,3),1),1);
    
    J = size(States.MidPointsFixedStates,2)-1;
    n = sum(SugarIdxR);
    FixedCov = zeros(n,J);
    for k=1:J
        FixedCov(:,k) = States.MidPointsFixedStates(States.FixedStatesMap(SugarAllStates(:,3),1+k),1+k);
    end
    Age   = double(States.AllStatesMap(SugarAllStates(:,1),1));
    
    OutReturn = AverageOut(r)*ones(n,1); % This value will be the same for everyone here
    
    X1 = [zeros(n,1) ones(n,1) zeros(n,1) Distance OutReturn zeros(n,J+1)];
    X2 = zeros(size(X1));
    X3 = [ones(n,1) zeros(n,1) Distance zeros(n,2) FixedCov Age];
    X = [X1 X2 X3];

    sumexp = sum(exp(X1*theta_fs.b_suc) + exp(X2*theta_fs.b_suc) + exp(X3*theta_fs.b_suc),2);

    prob_ini{r}(SugarIdxR,1) = exp(X1*theta_fs.b_suc)./sumexp;
    prob_ini{r}(SugarIdxR,2) = exp(X2*theta_fs.b_suc)./sumexp;
    prob_ini{r}(SugarIdxR,3) = exp(X3*theta_fs.b_suc)./sumexp;

    
    % Generate choice probabilities for non sugar states:
    
    NSugarIdx = (States.AllStatesMap(:,1)==1);
    NSugarAllStates = States.AllStatesMap(NSugarIdx,:);

    Distance = States.MidPointsFixedStates(States.FixedStatesMap(NSugarAllStates(:,3),1),1);
    J = size(States.MidPointsFixedStates,2)-1;
    n = sum(NSugarIdx);
    FixedCov = zeros(n,J);
    for k=1:J
        FixedCov(:,k) = States.MidPointsFixedStates(States.FixedStatesMap(NSugarAllStates(:,3),1+k),1+k);
    end
    Age   = double(States.AllStatesMap(NSugarAllStates(:,1),1));
    
    LUclass = double(States.FixedStatesMap(NSugarAllStates(:,3), end)); % Get Other Land Use class state
    LUclass = dummyvar(LUclass); % create dummy variables for clogit
    LUclass = LUclass(:,2:end); % clogit will be run with constant, so eliminate first column (not pasture nor cropland)
    % Sugarcane Road Distance (SRD)
    SRD     = States.SRDMidPoints(NSugarAllStates(:,4),1);
    OutReturn = AverageOut(r)*ones(n,1); % This value will be the same for everyone here

    % Covariates are action specific: plant sugar and do nothing: 
    X2 = [ones(n,1) zeros(n,1) Distance log(SRD) FixedCov LUclass];
    X1 = [zeros(n,1) OutReturn zeros(n,2+J+2)];
    X = [X1 X2];
    
    sumexp = sum(exp(X1*theta_fs.b_nsuc) + exp(X2*theta_fs.b_nsuc),2);

    prob_ini{r}(NSugarIdx,2) = exp(X1*theta_fs.b_nsuc)./sumexp;
    prob_ini{r}(NSugarIdx,3) = exp(X2*theta_fs.b_nsuc)./sumexp;

end




